Skip to content

Progress bars when sampling multiple chains #168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 25 commits into
base: main
Choose a base branch
from

Conversation

penelopeysm
Copy link
Member

@penelopeysm penelopeysm commented Jun 30, 2025

Summary

This PR introduces more informative progress bars with MCMCThreads which can be configured by the user (see options below).

In today's meeting it was generally agreed that having one progress bar per chain was more informative / useful, and that an upper bound could be placed to disable that by default if there were more than MAX_CHAINS_PROGRESS chains.

This PR sets MAX_CHAINS_PROGRESS = 10. The user can override the default on a case-by-case basis using the Symbol versions of progress, or can globally set the threshold using setmaxchainsprogress!(N).

For MCMCSerial no change is made (it already had one progress bar per chain).

For MCMCDistributed, progress = :overall is implemented (progress = true is a synonym for that) but progress = :perchain isn't, not for any good technical reason, but because it was just too much of a faff. I think that in principle it's possible to implement it similarly to how it's done for MCMCThreads (i.e., generate one UUID per chain), but I think each chain has to communicate back to the parent worker using a RemoteChannel, and there needs to be a way of disentangling the inputs from different chains (so you either need N RemoteChannels or the channel needs to take a chain ID). I suppose if somebody really, really wants it I could do it, but given that :overall already works I am hoping nobody will complain.

In a technical sense this is accomplished by passing richer objects in the progress keyword argument (see src/logging.jl). This does mean that people should not override progress in their own implementations (I have added a warning in the docs) but it is IMO better than using callbacks.

The videos below demonstrate the behaviour for MCMCThreads.

Closes #82
Closes TuringLang/Turing.jl#2264

Setup

using AbstractMCMC
struct M <: AbstractMCMC.AbstractModel end
struct S <: AbstractMCMC.AbstractSampler end
function AbstractMCMC.step(rng, ::M, ::S, state=nothing; kwargs...)
    sleep(0.001)
    rand(rng), nothing
end

progress=:perchain

# One progress bar per chain, plus one overall progress bar which tracks number of chains
sample(M(), S(), MCMCThreads(), 500, 6; progress=:perchain)
perchain.mov

progress=:overall

# One overall progress bar which tracks the number of samples
sample(M(), S(), MCMCThreads(), 500, 6; progress=:perchain)
overall.mov

progress=true

# Uses :perchain for 10 or fewer chains, :overall for more than 10 chains
# This is still the default option
sample(M(), S(), MCMCThreads(), 500, 6; progress=true)
sample(M(), S(), MCMCThreads(), 200, 11; progress=true)
true.mov

progress=false or progress=:none

# No progress bars at all
# `progress=:none` also works. We have to preserve the boolean values though so that
# `setprogress!(false)` works globally
sample(M(), S(), MCMCThreads(), 500, 6; progress=false)
false.mov

Bonus: :perchain and :overall on Pluto (it's pretty)

Screen.Recording.2025-06-30.at.23.51.20.mov

Bonus: :perchain in iJulia (in VSCode; it's not pretty)

Screen.Recording.2025-06-30.at.23.59.19.mov

Bonus: :overall in iJulia (works fine)

Screen.Recording.2025-06-30.at.23.59.46.mov

Copy link
Contributor

AbstractMCMC.jl documentation for PR #168 is available at:
https://TuringLang.github.io/AbstractMCMC.jl/previews/PR168/

@penelopeysm penelopeysm force-pushed the py/parallel-combined branch from 8a9d376 to 1195503 Compare June 30, 2025 22:31
@TuringLang TuringLang deleted a comment from github-actions bot Jun 30, 2025
@TuringLang TuringLang deleted a comment from github-actions bot Jun 30, 2025
@penelopeysm penelopeysm marked this pull request as ready for review July 1, 2025 13:58
@penelopeysm penelopeysm changed the title Progress bars Progress bars when sampling multiple chains Jul 1, 2025
@penelopeysm penelopeysm requested a review from nsiccha July 1, 2025 15:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Progress reporting in parallel sampling More useful logging when sampling multiple chains
1 participant